import asyncio
import json
import multiprocessing
import threading
from queue import Empty
from fastapi import FastAPI
from fastapi.websockets import WebSocket
import sys
from tts2 import create_recognizer, create_keyword_spotter, create_tts

from starlette.websockets import WebSocketDisconnect

try:
    import sounddevice as sd
except ImportError as e:
    print("Please install sounddevice first. You can use")
    print()
    print("  pip install sounddevice")
    print()
    print("to install it")
    sys.exit(-1)
from get_config import config
import logging

app = FastAPI()
audio_config = config["audio"]

print("启动服务中")
RATE = int(audio_config["rate"])  # 硬件采样率
TARGET_RATE = audio_config["target_rate"]  # 配置中的目标采样率
CHUNK = audio_config["chunk"]  # 每帧大小
print("WS服务启动完成")

async def send_message(websocket, msg):
    """异步发送 WebSocket 消息"""
    await websocket.send_text(json.dumps(msg,ensure_ascii=False))


def consumer_process(websocket_queue, stop_event):
    recognizer =create_recognizer()
    keyword_spotter = create_keyword_spotter()
    print("stt服务启动成功!")
    websocket_queue.put({"event": "ok"})
    stream = recognizer.create_stream()
    last_result = ""
    kws=""
    kws_stream = keyword_spotter.create_stream()
    with sd.InputStream(channels=1, dtype="float32", samplerate=RATE) as s:
        while not stop_event.is_set():
            samples, _ = s.read(CHUNK)
            samples = samples.reshape(-1)
            stream.accept_waveform(RATE, samples)
            while recognizer.is_ready(stream):
                recognizer.decode_stream(stream)
            kws_stream.accept_waveform(RATE, samples)
            while keyword_spotter.is_ready(kws_stream):
                 keyword_spotter.decode_stream(kws_stream)
            kws_result = keyword_spotter.get_result(kws_stream)
            if kws_result:
                kws=kws_result
            is_endpoint = recognizer.is_endpoint(stream)
            result = recognizer.get_result(stream)
            if result and is_endpoint:
                if kws:
                    msg = {"event": "endpoint", "text": result,"kws":kws}
                    kws=""
                else:
                    msg={"event": "endpoint", "text": result}
                # logging.info(msg)
                print(msg)
                websocket_queue.put(msg)
            if is_endpoint:
                recognizer.reset(stream)
        s.close()

def consumer_task(websocket_queue, stop_event, websocket,loop):
    """线程任务，异步从队列中获取数据并发送"""
    while not stop_event.is_set():
        try:
            # 异步从队列中获取消息
            msg = websocket_queue.get()  # 阻塞等待队列中的消息
            if msg:
                loop.create_task(send_message(websocket, msg))  # 异步发送消息
        except Empty:
            pass  # 队列为空，继续等待
        except Exception as e:
            logging.error(f"消费者线程发生错误: {e}")
            break
@app.websocket("/ws2")
async def useOnnx(websocket:WebSocket):
    await websocket.accept()
    await send_message(websocket, {"event": "start"})
    stop_event = multiprocessing.Event()
    websocket_queue = multiprocessing.Queue()
    loop = asyncio.get_event_loop()
    multiprocessing.set_start_method("spawn", force=True)  # 添加此行
    process = multiprocessing.Process(target=consumer_process, args=(websocket_queue, stop_event))
    process.start()
    consumer_thread = threading.Thread(target=consumer_task, args=(websocket_queue, stop_event, websocket, loop))
    consumer_thread.start()

    try:
        # 保持 WebSocket 打开状态
        while True:
            await websocket.receive_text()
    except WebSocketDisconnect:
        logging.info("WebSocket连接已断开")
    except Exception as e:
        logging.error(f"WebSocket发生错误: {e}")
    finally:
        # 停止消费者进程
        logging.info("释放资源...")
        stop_event.set()  # 设置停止标志
        process.join()  # 等待进程结束
        try:
            await websocket.close()
        except:
            pass
        logging.info("资源释放完成")


# 本地麦克风
@app.websocket('/tts/local')
async def useOnnxTts(websocket:WebSocket):
    await websocket.accept()
    stop_event = multiprocessing.Event()
    websocket_queue = multiprocessing.Queue()
    multiprocessing.set_start_method("spawn", force=True)
    from tts_server import tts_process
    process = multiprocessing.Process(target=tts_process, args=(websocket_queue, stop_event))
    process.start()
    try:
        # 保持 WebSocket 打开状态
        while True:
            receive_text=await websocket.receive_text()
            j=json.loads(receive_text)
            websocket_queue.put(j['text'])
    except WebSocketDisconnect:
        logging.info("WebSocket连接已断开")
    except Exception as e:
        logging.error(f"WebSocket发生错误: {e}")
    finally:
        websocket_queue.put("__shop")
        # 停止消费者进程
        logging.info("释放资源...")
        stop_event.set()  # 设置停止标志
        process.join()  # 等待进程结束
        try:
            await websocket.close()
        except:
            pass
        logging.info("资源释放完成")
